#include "red_black_tree.h"
#include "constants.h"

RBNode *GetMinimumNode(RBNode *node)
{
	while(node->left != NULL)
		node = node->left;

	return node;
}

RBNode *GetMaximumNode(RBNode *node)
{
	while(node->right != NULL)
		node = node->right;

	return node;
}

RBNode *TreeSuccessor(RBTree *tree, RBNode *node)
{
	RBNode *y;
	if(node->right != NULL)
		return GetMinimumNode(node->right);

	y = node->p;
	while(y != NULL && node == y->right)
	{
		node = y;
		y = y->p;
	}
	return y;
}

RBNode *TreePredecessor(RBTree *tree, RBNode *node)
{
	RBNode *y;

	if(node->left != NULL)
		return GetMaximumNode(node->left);

	y = node->p;
	while(y != NULL && node == y->left)
	{
		node = y;
		y = y->p;
	}
	return y;
}

RBNode *CreateNewRBNode(void *data)
{
	RBNode *node;

	node = (RBNode*)SafeMalloc(sizeof(RBNode));
	node->data = data;
	node->p = NULL;
	node->left = NULL;
	node->right = NULL;
	node->red = TRUE;

	return node;
}

void RotateLeft(RBTree *tree, RBNode *node)
{
	RBNode *tmp = node->right;

	//Set node's right child
	node->right = tmp->left;

	if(tmp->left)
		tmp->left->p = node;

	if(node->p)
	{
		tmp->p = node->p;
		node->p->left = tmp;
	}
	else
	{
		tree->root = tmp;
		tmp->p = NULL;
	}

	node->p = tmp;
	tmp->left = node;
}

void RotateRight(RBTree *tree, RBNode *node)
{
	RBNode *tmp = node->left;

	node->left = tmp->right;

	if(tmp->right)
		tmp->right->p = node;

	if(node->p)
	{
		tmp->p = node->p;
		node->p->right = tmp;
	}
	else
	{
		tree->root = tmp;
		tmp->p = NULL;
	}

	node->p = tmp;
	tmp->right = node;
}

void InsertFixup(RBTree *tree, RBNode *node)
{
	//If the node is the root node, change the node to black and return;
	if(node == tree->root)
	{
		node->red = FALSE;
		return;
	}

	//If the parent node is black, there is no need to fixup.
	if(node->p->red)
	{
		//If node's parent is the left one.
		if(node->p == node->p->p->left)
		{
			//If node's uncle node is red, just change the color and refixup.
			if(node->p->p->right != NULL && 
				node->p->p->right->red)
			{
				node->p->red = FALSE;
				node->p->p->right->red = FALSE;
				node->p->p->red = TRUE;
				InsertFixup(tree, node->p->p);
			}
			//If nodes' uncle node is black. Need rotate.
			else
			{
				if(node == node->p->right)
				{
					node = node->p;
					RotateLeft(tree, node);
				}
				node->p->red = FALSE;
				node->p->p->red = TRUE;
				RotateRight(tree, node->p->p);
			}
		}
		//Same as above, just make left and right reversed.
		else
		{
			if(node->p->p->left != NULL && 
				node->p->p->left->red)
			{
				node->p->red = FALSE;
				node->p->p->left->red = FALSE;
				node->p->p->red = TRUE;
				InsertFixup(tree, node->p->p);
			}
			else
			{
				if(node == node->p->left)
				{
					node = node->p;
					RotateRight(tree, node);
				}
				node->p->red = FALSE;
				node->p->p->red = TRUE;
				RotateLeft(tree, node->p->p);
			}
		}
	}
	tree->root->red = FALSE;
}

RBNode *RBNodeInsert(RBTree *tree, void *data)
{
	RBNode *pCur, *node;
	int compare;

	if(tree->root == NULL)
	{
		tree->root = CreateNewRBNode(data);
		tree->root->red = FALSE;
		tree->nodeCount = 1;
		return tree->root;
	}

	pCur = tree->root;
	while(pCur)
	{
		compare = tree->Compare(pCur->data, data);
		if(compare == 0)
			return pCur;

		if(compare < 0)
		{
			if(pCur->right != NULL)
				pCur = pCur->right;
			else
			{
				node = CreateNewRBNode(data);
				pCur->right = node;
				node->p = pCur;
				break;
			}
		}
		else
		{
			if(pCur->left != NULL)
				pCur = pCur->left;
			else
			{
				node = CreateNewRBNode(data);
				pCur->left = node;
				node->p = pCur;
				break;
			}
		}
	}
	tree->nodeCount++;
	InsertFixup(tree, node);
	return node;
}

RBTree* CreateRBTree(CompareFunc compareFunc)
{
	RBTree* tree = (RBTree*)SafeMalloc(sizeof(RBTree));
	tree->Compare = compareFunc;
	tree->root = NULL;
	tree->nodeCount = 0;
	return tree;
}

void ReleaseNode(RBNode *node)
{
	if(node)
	{
		ReleaseNode(node->left);
		ReleaseNode(node->right);
		free(node);
	}
}

void ReleaseTree(RBTree *tree)
{
	ReleaseNode(tree->root);
	free(tree);
}
